# database_tools.py
import pandas as pd
from typing import Dict, List
from langchain_core.tools import tool
from sqlalchemy import create_engine, inspect
import os

# --- 配置 ---
CSV_FILE_PATH = "data/sales_data.csv"  # 你的 CSV 文件路径
POSTGRES_URL = "postgresql+psycopg2://username:password@localhost:5432/analytics_db"  # 修改为你的连接信息

# 创建 PostgreSQL 引擎
engine = create_engine(POSTGRES_URL)


# --- 工具定义 ---

@tool
def get_csv_info() -> Dict:
    """Get information about the local CSV file: path, columns, and first few rows."""
    if not os.path.exists(CSV_FILE_PATH):
        return {"error": f"CSV file not found at {CSV_FILE_PATH}"}

    try:
        df = pd.read_csv(CSV_FILE_PATH)
        info = {
            "file_path": CSV_FILE_PATH,
            "columns": df.columns.tolist(),
            "num_rows": len(df),
            "sample_data": df.head().to_dict('records')  # 前5行
        }
        return info
    except Exception as e:
        return {"error": f"Error reading CSV: {str(e)}"}


@tool
def import_csv_to_postgres(table_name: str) -> Dict:
    """
    Import the local CSV file into a PostgreSQL table.
    The table name is specified. The table and schema are created automatically based on CSV structure.
    """
    if not os.path.exists(CSV_FILE_PATH):
        return {"error": f"CSV file not found at {CSV_FILE_PATH}"}

    try:
        df = pd.read_csv(CSV_FILE_PATH)

        # 使用 pandas 的 to_sql 自动创建表并插入数据
        # 如果表已存在，则替换
        df.to_sql(name=table_name, con=engine, if_exists='replace', index=False)

        # 验证导入
        inspector = inspect(engine)
        if table_name in inspector.get_table_names():
            row_count = engine.execute(f"SELECT COUNT(*) FROM {table_name}").scalar()
            return {
                "status": "success",
                "message": f"CSV data imported successfully into table '{table_name}'.",
                "rows_imported": row_count,
                "table_name": table_name
            }
        else:
            return {"error": f"Failed to create or find table '{table_name}' in PostgreSQL."}

    except Exception as e:
        return {"error": f"Error importing CSV to PostgreSQL: {str(e)}"}


@tool
def query_postgres(sql_query: str) -> Dict:
    """Execute a SQL query on the PostgreSQL database and return the results."""
    try:
        with engine.connect() as conn:
            result = conn.execute(sql_query)
            rows = result.fetchall()
            columns = result.keys()

            # 转换为字典列表
            results = [dict(zip(columns, row)) for row in rows]

            return {
                "query": sql_query,
                "results": results,
                "row_count": len(results)
            }
    except Exception as e:
        return {"error": f"SQL Error: {str(e)}"}


# 工具列表
tools = [get_csv_info, import_csv_to_postgres, query_postgres]